from typing import Any, Dict, Iterator, List, Mapping, Optional

from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.outputs import GenerationChunk
import asyncio

class CustomLLM(LLM):
    """自定义聊天模型，返回输入文本的前 `n` 个字符。

    在将实现贡献给 LangChain 时，应仔细记录模型的初始化参数，
    并提供初始化模型的示例。此外，如果有相关的底层模型文档或 API，
    也应提供相应链接。

    示例:

        .. code-block:: python

            model = CustomChatModel(n=2)
            result = model.invoke([HumanMessage(content="hello")])
            result = model.batch([[HumanMessage(content="hello")],
                                 [HumanMessage(content="world")]])
    """

    n: int
    """从输入文本的最后一条消息中返回的字符数。"""

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        """运行 LLM 以处理输入。

        需要重写此方法以实现 LLM 的逻辑。

        Args:
            prompt:  需要生成的输入文本。
            stop:  生成时使用的停止词，模型输出会在第一次遇到这些词时被截断。
                如果模型不支持停止词，建议抛出 NotImplementedError。
            run_manager: 运行时的回调管理器。
            **kwargs: 额外的关键字参数，通常用于传递给模型 API 调用。
        Returns:
            生成的文本字符串。注意，实际的生成结果**不应**包含输入文本。
        """
        if stop is not None:
            raise ValueError("不允许使用 stop 关键字参数。")
        return prompt[: self.n]

    def _stream(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[GenerationChunk]:
        """以流式方式运行 LLM 处理输入文本。

        支持流式输出的子类应重写此方法。

        如果未实现该方法，则 `stream` 调用会默认回退到普通调用，
        并将输出作为单个块返回。

        Args:
            prompt:  需要生成的输入文本。
            stop: 生成时使用的停止词，模型输出会在第一次遇到这些词时被截断。
            run_manager: 运行时的回调管理器。
            **kwargs: 额外的关键字参数，通常用于传递给模型 API 调用。

        Returns:
            生成文本的迭代器，每次返回一个 `GenerationChunk` 对象。
        """
        for char in prompt[: self.n]:
            chunk = GenerationChunk(text=char)
            if run_manager:
                run_manager.on_llm_new_token(chunk.text, chunk=chunk)

            yield chunk

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        """返回一个包含模型标识参数的字典。"""
        return {
            # 模型名运行用户在LLM监控应用程序中指定
            # 自定义tokens计数规则(e.g., 在LangSmith中，
            # 用户可以为他们的模型提供每个token的定价，
            # 并监控给定LLM的成本）
            "model_name": "CustomChatModel",
        }

    @property
    def _llm_type(self) -> str:
        """获取该聊天模型所使用的语言模型类型，仅用于日志记录。"""
        return "custom"
    
if __name__=="__main__":
    llm = CustomLLM(n=5)
    print(llm)
    print(llm.invoke("This is a foobar thing"))
    async def llmd():
        f = await llm.ainvoke("world")
        print(f)

    asyncio.run(llmd())
    print(llm.batch(["woof woof woof", "meow meow meow"]))